###############################################################
# R code template
#
# Association of Vitamin C Supplementation and Genetic
# Susceptibility with Multiple Sclerosis Risk:
# A Prospective Population-Based Cohort Study
#
# Andrea Nova
# University of Pavia
# andrea.nova@unipv.it
###############################################################

##Load libraries
library(cobalt)
library(rms)
library(coxed)
library(ipw)
library(survival)

## Dataset
# data contains:
#   VitaminC_Supp : Vitamin C supplementation (0 = no, 1 = yes)
#   TimeToMS      : Follow-up time
#   MS            : Incident multiple sclerosis (0/1)
#   MS_PRS        : Multiple sclerosis polygenic risk score
#   Confounders   : Covariates used for propensity-score estimation
#   PRS_Confounders : Sex, age, ethnicity, country of birth and
#                     first 10 genetic principal components

###############################################################
# 1. Inverse probability of treatment weighting
###############################################################

temp <- ipwpoint(
  exposure = VitaminC_Supp
  family = "binomial",
  link = "logit",
  numerator = ~ 1,
  denominator = ~ Confounders,
  data = data)
summary(temp$ipw.weights) #Check Max IPW < 10

#Save weights
data$IPW<-temp$ipw.weights

###############################################################
# 2. Assessment of covariate balance after weighting
###############################################################

#Check balance through standardized mean differences (SMDs)
love.plot(VitaminC_Supp ~ Confounders, data = data, weights = data$IPW, method = "weighting", binary = "std", var.order = "unadjusted",
          thresholds = c(m = 0.1), abs = FALSE)
#To obtain SMDs
bal.tab(VitaminC_Supp ~ Confounders, data = data, weights = data$IPW, method = "weighting")

###############################################################
# 3. Weighted Cox proportional hazards model
###############################################################

model <- coxph(Surv(TimeToMS, MS) ~ VitaminC_Supp, data = data, family=binomial, weights=data$IPW)

###############################################################
# 4. Non-linear interaction between vitamin C supplementation
#    and MS polygenic risk score
###############################################################
#MS-PRS modeled using restricted cubic splines with 3 knots.
library(rms)
model<-cph(
  Surv(TimeToMS, MS) ~ 
    VitaminC_Supp*rcs(MS_PRS,3)+PRS_Confounders,weights=data$IPW, data=data)
anova(model) #p-value for interaction and non-linear terms

#Obtain Hazard Ratios curve as a function of MS-PRS and RERI

#Range of MS-PRS values
prs_seq <- seq(
  quantile(data$MS_PRS, 0.001, na.rm=T),
  quantile(data$MS_PRS, 0.999, na.rm=T),
  length.out = 1000
)

#Counterfactual contrasts:
#Vitamin C supplement = 1, MS-PRS = each value in the range of prs_seq
newdata_vitC1 <- data.frame(
  MS_PRS = prs_seq,
  VitaminC_supp = 1
)

#Vitamin C supplement = 0, MS-PRS = each value in the range of prs_seq
newdata_vitC0 <- data.frame(
  MS_PRS = prs_seq,
  VitaminC_supp = 0
)

#Fix PRS confounders to the mean value
cov_means <- data.frame(lapply(data[PRS_Confounders], mean, na.rm = TRUE))

newdata_vitC1 <- cbind(newdata_vitC1, cov_means[rep(1, length(prs_seq)), ])
newdata_vitC0 <- cbind(newdata_vitC0, cov_means[rep(1, length(prs_seq)), ])

#Linear predictors based on the fitted model with the interaction term
lp1 <- predict(model, newdata = newdata_vitC1, type = "lp", se.fit = F)
lp0 <- predict(model, newdata = newdata_vitC0, type = "lp", se.fit = F)

#Counterfactual contrast: log(HR) of Vitamin C supplement use as a function of MS-PRS
logHR <- lp1 - lp0
HR_curve<-exp(logHR)

###############################################################
# 5. Additive interaction: Relative Risk due to Interaction (RERI).
###############################################################

#Exposures definition:

#Vitamin C supplement = 1, MS-PRS = 1
newdata_vitC11 <- data.frame(
  MS_PRS = 1,
  VitaminC_supp = 1
)
newdata_vitC11 <- cbind(newdata_vitC11, cov_means)
lp11 <- predict(model, newdata = newdata_vitC11, type = "lp", se.fit = F)

#Vitamin C supplement = 1, MS-PRS = -1
newdata_vitC10 <- data.frame(
  MS_PRS = -1,
  VitaminC_supp = 1
)
newdata_vitC10 <- cbind(newdata_vitC10, cov_means)
lp10 <- predict(model, newdata = newdata_vitC10, type = "lp", se.fit = F)

#Vitamin C supplement = 0, MS-PRS = 1
newdata_vitC01 <- data.frame(
  MS_PRS = 1,
  VitaminC_supp = 0
)
newdata_vitC01 <- cbind(newdata_vitC01, cov_means)
lp01 <- predict(model, newdata = newdata_vitC01, type = "lp", se.fit = F)

#Vitamin C supplement = 0, MS-PRS = -1
newdata_vitC00 <- data.frame(
  MS_PRS = -1,
  VitaminC_supp = 0
)
newdata_vitC00 <- cbind(newdata_vitC00, cov_means)
lp00 <- predict(model, newdata = newdata_vitC00, type = "lp", se.fit = F)

#Counterfacual contrast:
# RERI = HR(VitC=0,PRS=1) - HR(VitC=0,PRS=-1)
#      - HR(VitC=1,PRS=1) + HR(VitC=1,PRS=-1)

RERI_estimate<-exp(lp01) - exp(lp00)- exp(lp11)+ exp(lp10)


###############################################################
# 6. 95% Confidence intervals using Bias-corrected accelerated (BCa) method
###############################################################
# 1000 bootstrap samples were generated by resampling participants
# with replacement. For each sample, the interaction model was refitted,
# and predicted HR curves and RERI estimates were recalculated.
HR<-NULL
RERI<-NULL
for(i in 1:1000){
ALLsample_index<-1:nrow(data)
ALLsample_index<-sample(ALLsample_index,size=length(ALLsample_index),replace=T)

model<-cph(
    Surv(TimeToMS, MS) ~ 
    VitaminC_Supp*rcs(MS_PRS,3)+PRS_Confounders,data =data[ALLsample_index,],weights=data$IPW)


prs_seq <- seq(
  quantile(data$MS_PRS, 0.001, na.rm=T),
  quantile(data$MS_PRS, 0.999, na.rm=T),
  length.out = 1000
)

newdata_vitC1 <- data.frame(
  MS_PRS = prs_seq,
  VitaminC_supp = 1
)

newdata_vitC0 <- data.frame(
  MS_PRS = prs_seq,
  VitaminC_supp = 0
)

newdata_vitC1 <- cbind(newdata_vitC1, cov_means[rep(1, length(prs_seq)), ])
newdata_vitC0 <- cbind(newdata_vitC0, cov_means[rep(1, length(prs_seq)), ])

lp1 <- predict(model, newdata = newdata_vitC1, type = "lp", se.fit = F)
lp0 <- predict(model, newdata = newdata_vitC0, type = "lp", se.fit = F)

logHR <- lp1 - lp0

HR <- rbind(HR,exp(logHR))

newdata_vitC11 <- data.frame(
  MS_PRS = 1,
  VitaminC_supp = 1
)
newdata_vitC11 <- cbind(newdata_vitC11, cov_means)
lp11 <- predict(model, newdata = newdata_vitC11, type = "lp", se.fit = F)

newdata_vitC10 <- data.frame(
  MS_PRS = -1,
  VitaminC_supp = 1
)
newdata_vitC10 <- cbind(newdata_vitC10, cov_means)
lp10 <- predict(model, newdata = newdata_vitC10, type = "lp", se.fit = F)

newdata_vitC01 <- data.frame(
  MS_PRS = 1,
  VitaminC_supp = 0
)
newdata_vitC01 <- cbind(newdata_vitC01, cov_means)
lp01 <- predict(model, newdata = newdata_vitC01, type = "lp", se.fit = F)

newdata_vitC00 <- data.frame(
  MS_PRS = -1,
  VitaminC_supp = 0
)
newdata_vitC00 <- cbind(newdata_vitC00, cov_means)
lp00 <- predict(model, newdata = newdata_vitC00, type = "lp", se.fit = F)

RERI<-c(RERI,exp(lp01) - exp(lp00)- exp(lp11)+ exp(lp10))


print(i)
}

95% CIs for Hazard Ratios curve
library(coxed)
CLlower<-NULL
CLupper<-NULL
for(j in 1:ncol(HR)){
  CLlower<-c(CLlower,bca(log(HR[,j]))[1])
  CLupper<-c(CLupper,bca(log(HR[,j]))[2])
}

95% CIs for RERI
bca(RERI)


